{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HugeCTR Embedding  Plugin for TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook introduces a TensorFlow (TF) plugin for the HugeCTR embedding layer, embedding_plugin, where users may benefit from both the computational efficiency of the HugeCTR embedding layer and the ease of use of TensorFlow (TF)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is new ##\n",
    "- Support `Localized` embedding.\n",
    "- No need to split DNN model into two sub-models, which means embedding layer can be put inside the scope of MirroredStrategy. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check Docker Container ##\n",
    "Please make sure that you have started the notebook inside the running NGC docker container: `nvcr.io/nvstaging/merlin/merlin-tensorflow-training:0.5`. Several dynamic libraries have been installed to the system path `/usr/local/hugectr/lib/` that you'll have to load using TensorFlow. For convenient usage, you can directly import `hugectr_tf_ops_v2.py`, where we prepare the codes to load that dynamic library and wrap some operations, in your python script to be used with the embedding_plugin."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Verify Accuracy ##\n",
    "To verify whether the embedding_plugin can obtain correct result, you can generate synthetic data for testing purposes as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run this cell to clear all variables.\n",
    "%reset -f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow and some modules\n",
    "import tensorflow as tf\n",
    "# do not let TF allocate all GPU memory\n",
    "devices = tf.config.list_physical_devices(\"GPU\")\n",
    "for dev in devices:\n",
    "    tf.config.experimental.set_memory_growth(dev, True)\n",
    "    \n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO]: loadding from /workspace/hugectr/build/lib/libembedding_plugin_v2.so\n"
     ]
    }
   ],
   "source": [
    "# import hugectr_tf_ops.py to use embedding_plugin ops\n",
    "import sys\n",
    "sys.path.append(\"../tools/embedding_plugin/python/\")\n",
    "import hugectr_tf_ops_v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "init embedding table value:\n",
      " [[ 1.  2.  3.  4.]\n",
      " [ 5.  6.  7.  8.]\n",
      " [ 9. 10. 11. 12.]\n",
      " [13. 14. 15. 16.]\n",
      " [17. 18. 19. 20.]\n",
      " [21. 22. 23. 24.]\n",
      " [25. 26. 27. 28.]\n",
      " [29. 30. 31. 32.]]\n"
     ]
    }
   ],
   "source": [
    "# generate a random embedding table and show\n",
    "vocabulary_size = 8\n",
    "slot_num = 3\n",
    "embedding_vector_size = 4\n",
    "\n",
    "table = np.float32([i for i in range(1, vocabulary_size * embedding_vector_size + 1)]).reshape(vocabulary_size, embedding_vector_size)\n",
    "print(\"init embedding table value:\\n\", table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In HugeCTR, the corresponding dense shape of the input keys is `[batch_size, slot_num, max_nnz]`, and `0` is a valid key. Therefore, `-1` is used to denote invalid keys, which only occupy that position in the corresponding dense keys matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the dense shape of inputs keys: (4, 3, 2)\n"
     ]
    }
   ],
   "source": [
    "# generate random keys to lookup from embedding table.\n",
    "keys = np.array([[[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [2,  6]],  # nnz = 2\n",
    "                 \n",
    "                 [[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [-1, -1]], # nnz = 0\n",
    "                 \n",
    "                 [[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [6, -1]],  # nnz = 1\n",
    "                 \n",
    "                 [[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [2, -1]]], # nnz = 1\n",
    "                dtype=np.int64) \n",
    "print(\"the dense shape of inputs keys:\", keys.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distributed Embedding first forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [34. 36. 38. 40.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 0.  0.  0.  0.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [25. 26. 27. 28.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 9. 10. 11. 12.]]], shape=(4, 3, 4), dtype=float32) \n",
      "\n",
      "Localized Embedding first forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [34. 36. 38. 40.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 0.  0.  0.  0.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [25. 26. 27. 28.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 9. 10. 11. 12.]]], shape=(4, 3, 4), dtype=float32) \n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Distributed Embedding second forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [33.800995   35.800995   37.800995   39.800995  ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [ 0.          0.          0.          0.        ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [24.900497   25.900497   26.900497   27.900497  ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [ 8.900497    9.900497   10.900497   11.900497  ]]], shape=(4, 3, 4), dtype=float32) \n",
      "\n",
      "Localized Embedding second forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [33.800995   35.800995   37.800995   39.800995  ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [ 0.          0.          0.          0.        ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [24.900497   25.900497   26.900497   27.900497  ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [ 8.900497    9.900497   10.900497   11.900497  ]]], shape=(4, 3, 4), dtype=float32) \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# define a simple forward propagation and backward propagation with embedding_plugin\n",
    "# NOTE: cause hugectr_tf_ops_v2.init() can only be called once, \n",
    "# if you want to run this cell multi-times, please restart the kernel,\n",
    "# or explicitly release embedding_plugin resources by calling hugectr_tf_ops_v2.reset()\n",
    "\n",
    "# try release embedding plugin resources.\n",
    "hugectr_tf_ops_v2.reset()\n",
    "\n",
    "# hugectr_tf_ops embedding_plugin initialize\n",
    "hugectr_tf_ops_v2.init(visible_gpus=[0], seed=0, key_type='int64', value_type='float', batch_size=4, batch_size_eval=4)\n",
    "\n",
    "# create a distributed embedding_layer with embedding_plugin\n",
    "dis_embedding_name = hugectr_tf_ops_v2.create_embedding(init_value=table, opt_hparams=[0.1, 0.9, 0.99, 1e-3], \n",
    "                                          name_='embedding_verification', \n",
    "                                          max_vocabulary_size_per_gpu=vocabulary_size,\n",
    "                                          slot_num=slot_num, embedding_vec_size=embedding_vector_size,\n",
    "                                          embedding_type='distributed', max_nnz=2)\n",
    "\n",
    "# create a localized embedding_layer with embedding_plugin\n",
    "loc_embedding_name = hugectr_tf_ops_v2.create_embedding(init_value=table, opt_hparams=[0.1, 0.9, 0.99, 1e-3],\n",
    "                                          name_='embedding_verification',\n",
    "                                          max_vocabulary_size_per_gpu=vocabulary_size, \n",
    "                                          slot_num=slot_num, embedding_vec_size=embedding_vector_size,\n",
    "                                          embedding_type='localized', max_nnz=2, update_type='Global')\n",
    "\n",
    "# convert dense input keys to COO format\n",
    "reshape_keys = tf.reshape(keys, [-1, keys.shape[-1]])\n",
    "indices = tf.where(reshape_keys != -1)\n",
    "values = tf.gather_nd(reshape_keys, indices)\n",
    "row_indices = tf.transpose(indices, perm=[1, 0])\n",
    "\n",
    "# create a Variable used for backward propagation\n",
    "bp_trigger = tf.Variable(initial_value=1.0, trainable=True, dtype=tf.float32)\n",
    "\n",
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    tape.watch(bp_trigger)\n",
    "\n",
    "    # get distributed embedding forward result\n",
    "    dis_each_replicas = hugectr_tf_ops_v2.broadcast_then_convert_to_csr(dis_embedding_name, row_indices, values,\n",
    "                                                                        T = [tf.int32] * 1)\n",
    "    dis_forward_result = hugectr_tf_ops_v2.fprop(dis_embedding_name, 0, dis_each_replicas, bp_trigger, is_training=True)\n",
    "    print(\"Distributed Embedding first forward_result:\\n\", dis_forward_result, '\\n')\n",
    "\n",
    "    # get localized embedding forward result\n",
    "    loc_each_replicas = hugectr_tf_ops_v2.broadcast_then_convert_to_csr(loc_embedding_name, row_indices, values, \n",
    "                                                                       T = [tf.int32] * 1)\n",
    "    loc_forward_result = hugectr_tf_ops_v2.fprop(loc_embedding_name, 0, loc_each_replicas, bp_trigger, is_training=True)\n",
    "    print(\"Localized Embedding first forward_result:\\n\", loc_forward_result, '\\n')\n",
    "\n",
    "# compute gradients & update params\n",
    "dis_grads = tape.gradient(dis_forward_result, bp_trigger)\n",
    "loc_grads = tape.gradient(loc_forward_result, bp_trigger)\n",
    "    \n",
    "# do second forward propagation to check whether embedding table is updated.\n",
    "dis_forward_result_2 = hugectr_tf_ops_v2.fprop(dis_embedding_name, 0, dis_each_replicas, bp_trigger, is_training=True)\n",
    "loc_forward_result_2 = hugectr_tf_ops_v2.fprop(loc_embedding_name, 0, loc_each_replicas, bp_trigger, is_training=True)\n",
    "print(\"-\"*100)\n",
    "print(\"Distributed Embedding second forward_result:\\n\", dis_forward_result_2, '\\n')\n",
    "print(\"Localized Embedding second forward_result:\\n\", loc_forward_result_2, '\\n')\n",
    "\n",
    "# explicitly release embedding plugin resources\n",
    "hugectr_tf_ops_v2.reset()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [34. 36. 38. 40.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 0.  0.  0.  0.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [25. 26. 27. 28.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 9. 10. 11. 12.]]], shape=(4, 3, 4), dtype=float32)\n",
      "\n",
      "\n",
      "tf second forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [33.800995   35.800995   37.800995   39.800995  ]]\n",
      "\n",
      " [[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [ 0.          0.          0.          0.        ]]\n",
      "\n",
      " [[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [24.900497   25.900497   26.900497   27.900497  ]]\n",
      "\n",
      " [[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [ 8.900497    9.900497   10.900497   11.900497  ]]], shape=(4, 3, 4), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# similarly, use original tensorflow op to compare whether results are consistent.\n",
    "\n",
    "# define a tf embedding layer\n",
    "class EmbeddingLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, vocabulary_size, embedding_vec_size,\n",
    "                init_value):\n",
    "        super(EmbeddingLayer, self).__init__()\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size\n",
    "        self.init_value = init_value\n",
    "        \n",
    "    def build(self, _):\n",
    "        self.Var = self.add_weight(shape=(self.vocabulary_size, self.embedding_vec_size),\n",
    "                                         initializer=tf.constant_initializer(value=self.init_value))\n",
    "        \n",
    "    def call(self, inputs):\n",
    "        return tf.nn.embedding_lookup_sparse(self.Var, inputs, sp_weights=None, combiner=\"sum\")\n",
    "    \n",
    "with tf.GradientTape() as tape:\n",
    "    # reshape keys into [batch_size * slot_num, max_nnz]\n",
    "    reshape_keys = np.reshape(keys, newshape=(-1, keys.shape[-1]))\n",
    "    indices = tf.where(reshape_keys != -1)\n",
    "    values = tf.gather_nd(reshape_keys, indices)\n",
    "\n",
    "    # define a layer\n",
    "    tf_layer = EmbeddingLayer(vocabulary_size, embedding_vector_size, table)\n",
    "    \n",
    "    # wrap input keys components into a SparseTensor\n",
    "    sparse_tensor = tf.sparse.SparseTensor(indices, values, reshape_keys.shape)\n",
    "    \n",
    "    tf_forward = tf_layer(sparse_tensor)\n",
    "    print(\"tf forward_result:\\n\", tf.reshape(tf_forward, [keys.shape[0], keys.shape[1], tf_forward.shape[-1]]))\n",
    "    \n",
    "    # define an optimizer\n",
    "    optimizer = tf.keras.optimizers.Adam(learning_rate=0.1, beta_1=0.9, beta_2=0.99, epsilon=1e-3)\n",
    "    \n",
    "    # compute gradients & update params\n",
    "    grads = tape.gradient(tf_forward, tf_layer.trainable_weights)\n",
    "    optimizer.apply_gradients(zip(grads, tf_layer.trainable_weights))\n",
    "    \n",
    "    # do second forward propagation to check whether params are updated.\n",
    "    tf_forward_2 = tf_layer(sparse_tensor)\n",
    "    print(\"\\n\")\n",
    "    print(\"tf second forward_result:\\n\", tf.reshape(tf_forward_2, [keys.shape[0], keys.shape[1], tf_forward_2.shape[-1]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Consistent in first forward propagation for both Distributed & Localized Embedding? True\n",
      "Consistent in second forward propagation for both Distributed & Localized Embedding? True\n"
     ]
    }
   ],
   "source": [
    "# assert whether embedding_plugin's results are consistent with tensorflow original ops\n",
    "\n",
    "# verify first forward results consistency\n",
    "dis_first_forward_consistent = np.allclose(dis_forward_result.numpy(), \n",
    "                                tf.reshape(tf_forward, [keys.shape[0], keys.shape[1], tf_forward.shape[-1]]).numpy())\n",
    "loc_first_forward_consistent = np.allclose(loc_forward_result.numpy(),\n",
    "                                tf.reshape(tf_forward, [keys.shape[0], keys.shape[1], tf_forward.shape[-1]]).numpy())\n",
    "print(\"Consistent in first forward propagation for both Distributed & Localized Embedding?\", \n",
    "     (dis_first_forward_consistent and loc_first_forward_consistent))\n",
    "\n",
    "# verify second forward results consistency\n",
    "dis_second_forward_consistent = np.allclose(dis_forward_result_2.numpy(), \n",
    "                                tf.reshape(tf_forward_2, [keys.shape[0], keys.shape[1], tf_forward_2.shape[-1]]))\n",
    "loc_second_forward_consistent = np.allclose(loc_forward_result_2.numpy(),\n",
    "                                tf.reshape(tf_forward_2, [keys.shape[0], keys.shape[1], tf_forward_2.shape[-1]]))\n",
    "print(\"Consistent in second forward propagation for both Distributed & Localized Embedding?\", \n",
    "      (dis_second_forward_consistent and loc_second_forward_consistent))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results from embedding_plugins and original TF ops are consistent in both first and second forward propagation for both `Distributed Embedding` and `Localized Embedding`, which means the embedding_plugin can get the same forward result and perform the same backward propagation as TF ops. Therefore, the embedding_plugin can obtain correct results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DeepFM demo ##\n",
    "In this notebook, TF 2.x is used to build the DeepFM model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Models with the Embedding_Plugin ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first, import tensorflow and import plugin ops from hugectr_tf_ops_v2.py\n",
    "import tensorflow as tf\n",
    "# do not let TF allocate all GPU memory\n",
    "devices = tf.config.list_physical_devices(\"GPU\")\n",
    "for dev in devices:\n",
    "    tf.config.experimental.set_memory_growth(dev, True)\n",
    "import sys\n",
    "sys.path.append(\"../tools/embedding_plugin/python/\")\n",
    "import hugectr_tf_ops_v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define TF layers\n",
    "class Multiply(tf.keras.layers.Layer):\n",
    "    def __init__(self, out_units):\n",
    "        super(Multiply, self).__init__()\n",
    "        self.out_units = out_units\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        self.w = self.add_weight(name='weight_vector', shape=(input_shape[1], self.out_units),\n",
    "                                 initializer='glorot_uniform', trainable=True)\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        return inputs * self.w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build DeepFM with plugin ops\n",
    "class DeepFM_PluginEmbedding(tf.keras.models.Model):\n",
    "    def __init__(self, \n",
    "                 vocabulary_size, \n",
    "                 embedding_vec_size,\n",
    "                 dropout_rate, # list of float\n",
    "                 deep_layers, # list of int\n",
    "                 initializer,\n",
    "                 gpus,\n",
    "                 batch_size,\n",
    "                 batch_size_eval,\n",
    "                 embedding_type = 'localized',\n",
    "                 slot_num=1,\n",
    "                 seed=123):\n",
    "        super(DeepFM_PluginEmbedding, self).__init__()\n",
    "        tf.keras.backend.clear_session()\n",
    "        tf.compat.v1.set_random_seed(seed)\n",
    "\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.deep_layers = deep_layers\n",
    "        self.gpus = gpus\n",
    "        self.batch_size = batch_size\n",
    "        self.batch_size_eval = batch_size_eval \n",
    "        self.slot_num = slot_num\n",
    "        self.embedding_type = embedding_type\n",
    "\n",
    "        if isinstance(initializer, str):\n",
    "            initializer = False\n",
    "            \n",
    "        # when building model with embedding_plugin ops, init() should be called prior to any other ops.\n",
    "        hugectr_tf_ops_v2.init(visible_gpus=gpus, seed=seed, key_type='int64', value_type='float', \n",
    "                        batch_size=batch_size, batch_size_eval=batch_size_eval)\n",
    "        \n",
    "        # create a embedding_plugin layer\n",
    "        self.embedding_name = hugectr_tf_ops_v2.create_embedding(init_value=initializer, name_='hugectr_embedding',\n",
    "                                            embedding_type=embedding_type, optimizer_type='Adam',\n",
    "                                            max_vocabulary_size_per_gpu=(self.vocabulary_size // len(self.gpus)) + 1,\n",
    "                                            opt_hparams=[0.1, 0.9, 0.99, 1e-5], update_type='Local',\n",
    "                                            atomic_update=True, scaler=1.0, slot_num=self.slot_num, \n",
    "                                            max_nnz=1, max_feature_num=1*self.slot_num, \n",
    "                                            embedding_vec_size=self.embedding_vec_size + 1, combiner='sum')\n",
    "        \n",
    "        # other layers with TF original ops\n",
    "        self.deep_dense = []\n",
    "        for i, deep_units in enumerate(self.deep_layers):\n",
    "            self.deep_dense.append(tf.keras.layers.Dense(units=deep_units, activation=None, use_bias=True,\n",
    "                                                         kernel_initializer='glorot_normal', \n",
    "                                                         bias_initializer='glorot_normal'))\n",
    "            self.deep_dense.append(tf.keras.layers.Dropout(dropout_rate[i]))\n",
    "        self.deep_dense.append(tf.keras.layers.Dense(units=1, activation=None, use_bias=True,\n",
    "                                                     kernel_initializer='glorot_normal',\n",
    "                                                     bias_initializer=tf.constant_initializer(0.01)))\n",
    "        self.add_layer = tf.keras.layers.Add()\n",
    "        self.y_act = tf.keras.layers.Activation(activation='sigmoid')\n",
    "\n",
    "        self.dense_multi = Multiply(1)\n",
    "        self.dense_embedding = Multiply(self.embedding_vec_size)\n",
    "\n",
    "        self.concat_1 = tf.keras.layers.Concatenate()\n",
    "        self.concat_2 = tf.keras.layers.Concatenate()\n",
    "        \n",
    "    def build(self, _):\n",
    "        self.bp_trigger = self.add_weight(name='bp_trigger', shape=(1,), dtype=tf.float32, trainable=True)\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, dense_feature, each_replica, training=True):\n",
    "        \"\"\"\n",
    "        forward propagation.\n",
    "        #arguments:\n",
    "            dense_feature: [batch_size, dense_dim]\n",
    "        \"\"\"\n",
    "        with tf.name_scope(\"embedding_and_slice\"):\n",
    "            dense_0 = tf.cast(tf.expand_dims(dense_feature, 2), dtype=tf.float32) # [batchsize, dense_dim, 1]\n",
    "            dense_mul = self.dense_multi(dense_0) # [batchsize, dense_dim, 1]\n",
    "            dense_emb = self.dense_embedding(dense_0) # [batchsize, dense_dim, embedding_vec_size]\n",
    "            dense_mul = tf.reshape(dense_mul, [dense_mul.shape[0], -1]) # [batchsize, dense_dim * 1]\n",
    "            dense_emb = tf.reshape(dense_emb, [dense_emb.shape[0], -1]) # [batchsize, dense_dim * embedding_vec_size]\n",
    "\n",
    "            sparse = hugectr_tf_ops_v2.fprop(self.embedding_name, 0, #replica_ctx.replica_id_in_sync_group,\n",
    "                                                each_replica, self.bp_trigger, is_training=training) # [batch_size, self.slot_num, self.embedding_vec_size + 1]\n",
    "            \n",
    "            sparse_1 = tf.slice(sparse, [0, 0, self.embedding_vec_size], [-1, self.slot_num, 1]) #[batchsize, slot_num, 1]\n",
    "            sparse_1 = tf.squeeze(sparse_1, 2) # [batchsize, slot_num]\n",
    "\n",
    "            sparse_emb = tf.slice(sparse, [0, 0, 0], [-1, self.slot_num, self.embedding_vec_size]) #[batchsize, slot_num, embedding_vec_size]\n",
    "            sparse_emb = tf.reshape(sparse_emb, [-1, self.slot_num * self.embedding_vec_size]) #[batchsize, slot_num * embedding_vec_size]\n",
    "        \n",
    "        with tf.name_scope(\"FM\"):\n",
    "            with tf.name_scope(\"first_order\"):\n",
    "                first = self.concat_1([dense_mul, sparse_1]) # [batchsize, dense_dim + slot_num]\n",
    "                first_out = tf.reduce_sum(first, axis=-1, keepdims=True) # [batchsize, 1]\n",
    "                \n",
    "            with tf.name_scope(\"second_order\"):\n",
    "                hidden = self.concat_2([dense_emb, sparse_emb]) # [batchsize, (dense_dim + slot_num) * embedding_vec_size]\n",
    "                second = tf.reshape(hidden, [-1, dense_feature.shape[1] + self.slot_num, self.embedding_vec_size])\n",
    "                square_sum = tf.math.square(tf.math.reduce_sum(second, axis=1, keepdims=True)) # [batchsize, 1, embedding_vec_size]\n",
    "                sum_square = tf.math.reduce_sum(tf.math.square(second), axis=1, keepdims=True) # [batchsize, 1, embedding_vec_size]\n",
    "                \n",
    "                second_out = 0.5 * (sum_square - square_sum) # [batchsize, 1, embedding_vec_size]\n",
    "                second_out = tf.math.reduce_sum(second_out, axis=-1, keepdims=False) # [batchsize, 1]\n",
    "                \n",
    "        with tf.name_scope(\"Deep\"):\n",
    "            for i, layer in enumerate(self.deep_dense):\n",
    "                if i % 2 == 0: # dense\n",
    "                    hidden = layer(hidden)\n",
    "                else: # dropout\n",
    "                    hidden = layer(hidden, training)\n",
    "\n",
    "        y = self.add_layer([hidden, first_out, second_out])\n",
    "        y = self.y_act(y) # [batchsize, 1]\n",
    "\n",
    "        return y\n",
    "    \n",
    "    @property\n",
    "    def get_embedding_name(self):\n",
    "        return self.embedding_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above cells use embedding plugin ops and TF layers to define a TF DeepFM model. Similarly, define an embedding layer with TF original ops, and define a DeepFM model with that layer. Because embedding_plugin supports model parallelism, the parameters of the original TF embedding layer are equally distributed to each GPU for a fair performance comparison."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Models with the Original TF Ops ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a TF embedding layer with TF original ops\n",
    "class OriginalEmbedding(tf.keras.layers.Layer):\n",
    "    def __init__(self, \n",
    "                 vocabulary_size,\n",
    "                 embedding_vec_size,\n",
    "                 initializer='uniform',\n",
    "                 combiner=\"sum\",\n",
    "                 gpus=[0]):\n",
    "        super(OriginalEmbedding, self).__init__()\n",
    "\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size \n",
    "        if isinstance(initializer, str):\n",
    "            self.initializer = tf.keras.initializers.get(initializer)\n",
    "        else:\n",
    "            self.initializer = initializer\n",
    "        if combiner not in [\"sum\", \"mean\"]:\n",
    "            raise RuntimeError(\"combiner must be one of \\{'sum', 'mean'\\}.\")\n",
    "        self.combiner = combiner\n",
    "        if (not isinstance(gpus, list)) and (not isinstance(gpus, tuple)):\n",
    "            raise RuntimeError(\"gpus must be a list or tuple.\")\n",
    "        self.gpus = gpus\n",
    "\n",
    "    def build(self, _):\n",
    "        if isinstance(self.initializer, tf.keras.initializers.Initializer):\n",
    "            if len(self.gpus) > 1:\n",
    "                self.embeddings_params = list()\n",
    "                mod_size = self.vocabulary_size % len(self.gpus)\n",
    "                vocabulary_size_each_gpu = [(self.vocabulary_size // len(self.gpus)) + (1 if dev_id < mod_size else 0)\n",
    "                                            for dev_id in range(len(self.gpus))]\n",
    "\n",
    "                for i, gpu in enumerate(self.gpus):\n",
    "                    with tf.device(\"/gpu:%d\" %gpu):\n",
    "                        params_i = self.add_weight(name=\"embedding_\" + str(gpu), \n",
    "                                                   shape=(vocabulary_size_each_gpu[i], self.embedding_vec_size),\n",
    "                                                   initializer=self.initializer)\n",
    "                    self.embeddings_params.append(params_i)\n",
    "\n",
    "            else:\n",
    "                self.embeddings_params = self.add_weight(name='embeddings', \n",
    "                                                        shape=(self.vocabulary_size, self.embedding_vec_size),\n",
    "                                                        initializer=self.initializer)\n",
    "        else:\n",
    "            self.embeddings_params = self.initializer\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, keys, output_shape):\n",
    "        result = tf.nn.embedding_lookup_sparse(self.embeddings_params, keys, \n",
    "                                             sp_weights=None, combiner=self.combiner)\n",
    "        return tf.reshape(result, output_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define DeepFM model with original TF embedding layer\n",
    "class DeepFM_OriginalEmbedding(tf.keras.models.Model):\n",
    "    def __init__(self, \n",
    "                 vocabulary_size, \n",
    "                 embedding_vec_size,\n",
    "                 dropout_rate, # list of float\n",
    "                 deep_layers, # list of int\n",
    "                 initializer,\n",
    "                 gpus,\n",
    "                 batch_size,\n",
    "                 batch_size_eval,\n",
    "                 embedding_type = 'localized',\n",
    "                 slot_num=1,\n",
    "                 seed=123):\n",
    "        super(DeepFM_OriginalEmbedding, self).__init__()\n",
    "        tf.keras.backend.clear_session()\n",
    "        tf.compat.v1.set_random_seed(seed)\n",
    "\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.deep_layers = deep_layers\n",
    "        self.gpus = gpus\n",
    "        self.batch_size = batch_size\n",
    "        self.batch_size_eval = batch_size_eval \n",
    "        self.slot_num = slot_num\n",
    "        self.embedding_type = embedding_type\n",
    "\n",
    "        self.original_embedding_layer = OriginalEmbedding(vocabulary_size=vocabulary_size, \n",
    "                                            embedding_vec_size=embedding_vec_size + 1, \n",
    "                                            initializer=initializer, gpus=gpus)\n",
    "        self.deep_dense = []\n",
    "        for i, deep_units in enumerate(self.deep_layers):\n",
    "            self.deep_dense.append(tf.keras.layers.Dense(units=deep_units, activation=None, use_bias=True,\n",
    "                                                         kernel_initializer='glorot_normal', \n",
    "                                                         bias_initializer='glorot_normal'))\n",
    "            self.deep_dense.append(tf.keras.layers.Dropout(dropout_rate[i]))\n",
    "        self.deep_dense.append(tf.keras.layers.Dense(units=1, activation=None, use_bias=True,\n",
    "                                                     kernel_initializer='glorot_normal',\n",
    "                                                     bias_initializer=tf.constant_initializer(0.01)))\n",
    "        self.add_layer = tf.keras.layers.Add()\n",
    "        self.y_act = tf.keras.layers.Activation(activation='sigmoid')\n",
    "\n",
    "        self.dense_multi = Multiply(1)\n",
    "        self.dense_embedding = Multiply(self.embedding_vec_size)\n",
    "\n",
    "        self.concat_1 = tf.keras.layers.Concatenate()\n",
    "        self.concat_2 = tf.keras.layers.Concatenate()\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, dense_feature, sparse_feature, training=True):\n",
    "        \"\"\"\n",
    "        forward propagation.\n",
    "        #arguments:\n",
    "            dense_feature: [batch_size, dense_dim]\n",
    "            sparse_feature: for OriginalEmbedding, it is a SparseTensor, and the dense shape is [batch_size * slot_num, max_nnz];\n",
    "                            for PluginEmbedding, it is a list of [row_offsets, value_tensors, nnz_array]. \n",
    "        \"\"\"\n",
    "        with tf.name_scope(\"embedding_and_slice\"):\n",
    "            dense_0 = tf.cast(tf.expand_dims(dense_feature, 2), dtype=tf.float32) # [batchsize, dense_dim, 1]\n",
    "            dense_mul = self.dense_multi(dense_0) # [batchsize, dense_dim, 1]\n",
    "            dense_emb = self.dense_embedding(dense_0) # [batchsize, dense_dim, embedding_vec_size]\n",
    "            dense_mul = tf.reshape(dense_mul, [dense_mul.shape[0], -1]) # [batchsize, dense_dim * 1]\n",
    "            dense_emb = tf.reshape(dense_emb, [dense_emb.shape[0], -1]) # [batchsize, dense_dim * embedding_vec_size]\n",
    "\n",
    "            sparse = self.original_embedding_layer(sparse_feature, output_shape=[-1, self.slot_num, self.embedding_vec_size + 1])\n",
    "\n",
    "            sparse_1 = tf.slice(sparse, [0, 0, self.embedding_vec_size], [-1, self.slot_num, 1]) #[batchsize, slot_num, 1]\n",
    "            sparse_1 = tf.squeeze(sparse_1, 2) # [batchsize, slot_num]\n",
    "\n",
    "            sparse_emb = tf.slice(sparse, [0, 0, 0], [-1, self.slot_num, self.embedding_vec_size]) #[batchsize, slot_num, embedding_vec_size]\n",
    "            sparse_emb = tf.reshape(sparse_emb, [-1, self.slot_num * self.embedding_vec_size]) #[batchsize, slot_num * embedding_vec_size]\n",
    "        \n",
    "        with tf.name_scope(\"FM\"):\n",
    "            with tf.name_scope(\"first_order\"):\n",
    "                first = self.concat_1([dense_mul, sparse_1]) # [batchsize, dense_dim + slot_num]\n",
    "                first_out = tf.reduce_sum(first, axis=-1, keepdims=True) # [batchsize, 1]\n",
    "                \n",
    "            with tf.name_scope(\"second_order\"):\n",
    "                hidden = self.concat_2([dense_emb, sparse_emb]) # [batchsize, (dense_dim + slot_num) * embedding_vec_size]\n",
    "                second = tf.reshape(hidden, [-1, dense_feature.shape[1] + self.slot_num, self.embedding_vec_size])\n",
    "                square_sum = tf.math.square(tf.math.reduce_sum(second, axis=1, keepdims=True)) # [batchsize, 1, embedding_vec_size]\n",
    "                sum_square = tf.math.reduce_sum(tf.math.square(second), axis=1, keepdims=True) # [batchsize, 1, embedding_vec_size]\n",
    "                \n",
    "                second_out = 0.5 * (sum_square - square_sum) # [batchsize, 1, embedding_vec_size]\n",
    "                second_out = tf.math.reduce_sum(second_out, axis=-1, keepdims=False) # [batchsize, 1]\n",
    "                \n",
    "        with tf.name_scope(\"Deep\"):\n",
    "            for i, layer in enumerate(self.deep_dense):\n",
    "                if i % 2 == 0: # dense\n",
    "                    hidden = layer(hidden)\n",
    "                else: # dropout\n",
    "                    hidden = layer(hidden, training)\n",
    "\n",
    "        y = self.add_layer([hidden, first_out, second_out])\n",
    "        y = self.y_act(y) # [batchsize, 1]\n",
    "\n",
    "        return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dataset is needed to use these models for training. [Kaggle Criteo datasets](http://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset/) provided by CriteoLabs is used as the training dataset. The original training set contains 45,840,617 examples. Each example contains a label (0 by default or 1 if the ad was clicked) and 39 features in which 13 of them are integer and the other 26 are categorial. Since TFRecord is suitable for the training process and the Criteo dataset is missing numerous values across the feature columns, preprocessing is needed. The original test set won't be used because it doesn't contain labels."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset processing ###\n",
    "1. Download dataset from [https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/](http://azuremlsampleexperiments.blob.core.windows.net/criteo/day_1.gz).\n",
    "2. Extract the dataset by running the following command. \n",
    "    ```shell\n",
    "    $ gunzip day_1.gz\n",
    "    ```\n",
    "    \n",
    "3. The whole dataset is too large, so get a subset with\n",
    "    ```shell\n",
    "    $ head -n 45840617 day_1 > train.txt\n",
    "    ```\n",
    "4. Preprocess the datast and set missing values.\n",
    "Preprocessing functions are defined in [preprocess.py](../tools/embedding_plugin/performance_profile/preprocess.py). Open that file and check the codes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify source csv name and output csv name, run this command will do the preprocessing.\n",
    "# Warning: this command will take serveral hours to do preprocessing.\n",
    "%run ../tools/embedding_plugin/performance_profile/preprocess.py \\\n",
    "    --src_csv_path=../tools/embedding_plugin/train.txt \\\n",
    "    --dst_csv_path=../tools/embedding_plugin/train.out.txt \\\n",
    "    --normalize_dense=0 --feature_cross=0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. Split the dataset by running the following commands:\n",
    "```shell\n",
    "$ head -n 36672493 train.out.txt > train\n",
    "$ tail -n 9168124 train.out.txt > valtest\n",
    "$ head -n 4584062 valtest > val\n",
    "$ tail -n 4584062 valtest > test\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. Convert the dataset to a TFRecord file. Converting functions are defined in [txt2tfrecord.py](../tools/embedding_plugin/performance_profile/txt2tfrecord.py). Open that file and check the codes.\n",
    "After the data preprocessing is completed, *.tfrecord file(s) will be generated, which can be used for training. The training loop can now be configured to use the dataset and models to perform the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify source name and output tfrecord name, run this command will do the converting.\n",
    "# Warning: this command will take half an hour to do converting.\n",
    "%run ../tools/embedding_plugin/performance_profile/txt2tfrecord.py \\\n",
    "    --src_txt_name=train \\\n",
    "    --dst_tfrecord_name=train.tfrecord \\\n",
    "    --normalized=0 --use_multi_process=1 \\\n",
    "    --shard_num=1 \n",
    "    # if multi tfrecord files are wanted, set shard_num to the number of files."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define training loop and do training ###\n",
    "In [read_data.py](../tools/embedding_plugin/performance_profile/read_data.py), some preprocessing and TF data reading pipeline creation functions are defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO]: loadding from /workspace/hugectr/build/lib/libembedding_plugin.so\n"
     ]
    }
   ],
   "source": [
    "# set env path, so that some modules can be imported\n",
    "sys.path.append(\"../tools/embedding_plugin/performance_profile/\")\n",
    "\n",
    "import txt2tfrecord as utils\n",
    "from read_data import CreateDataset\n",
    "import time\n",
    "import logging\n",
    "logging.basicConfig(format='%(asctime)s %(message)s')\n",
    "logging.root.setLevel('INFO')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose wich model for training\n",
    "which_model = \"Plugin\" # change it to \"Original\", if you want to try the model define with original tf ops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set some hyper parameters for training process\n",
    "if (\"Plugin\" == which_model):\n",
    "    batch_size = 16384\n",
    "    n_epochs = 1\n",
    "    distribute_keys = 1 \n",
    "    gpus = [0] # use GPU0\n",
    "    embedding_type = 'distributed'\n",
    "    vocabulary_size = 1737710\n",
    "    embedding_vec_size = 10\n",
    "    slot_num = 26\n",
    "    batch_size_eval = 1 * len(gpus)\n",
    "    \n",
    "elif (\"Original\" == which_model):\n",
    "    batch_size = 16384\n",
    "    n_epochs = 1\n",
    "    distribute_keys = 0\n",
    "    gpus = [0] # use GPU0\n",
    "    vocabulary_size = 1737710\n",
    "    embedding_vec_size = 10\n",
    "    slot_num = 26\n",
    "    batch_size_eval = 1 * len(gpus)\n",
    "    embedding_type = 'distributed'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define feature_description to read tfrecord examples.\n",
    "cols = [utils.idx2key(idx, False) for idx in range(0, utils.NUM_TOTAL_COLUMNS)]\n",
    "feature_desc = dict()\n",
    "for col in cols:\n",
    "    if col == 'label' or col.startswith(\"I\"):\n",
    "        feature_desc[col] = tf.io.FixedLenFeature([], tf.int64) # scaler\n",
    "    else: \n",
    "        feature_desc[col] = tf.io.FixedLenFeature([1], tf.int64) # [slot_num, nnz]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# please set data_path to your tfrecord\n",
    "data_path = \"../tools/embedding_plugin/performance_profile/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create tfrecord reading pipeling\n",
    "dataset_names = [data_path + \"./train.tfrecord\"]\n",
    "dataset = CreateDataset(dataset_names=dataset_names,\n",
    "                       feature_desc=feature_desc,\n",
    "                       batch_size=batch_size,\n",
    "                       n_epochs=n_epochs,\n",
    "                       slot_num=slot_num,\n",
    "                       max_nnz=1,\n",
    "                       convert_to_csr=False,\n",
    "                       gpu_count=len(gpus),\n",
    "                       embedding_type=embedding_type,\n",
    "                       get_row_indices=True)()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define loss function and optimizer used in other TF layers.\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)\n",
    "loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create model instance\n",
    "if \"Original\" == which_model:\n",
    "    model = DeepFM_OriginalEmbedding(vocabulary_size=vocabulary_size, embedding_vec_size=embedding_vec_size, \n",
    "                                     embedding_type=embedding_type,\n",
    "                                     dropout_rate=[0.5] * 10, deep_layers=[1024] * 10,\n",
    "                                     initializer='uniform', gpus=gpus, batch_size=batch_size, \n",
    "                                     batch_size_eval=batch_size_eval,\n",
    "                                     slot_num=slot_num)\n",
    "elif \"Plugin\" == which_model:\n",
    "    hugectr_tf_ops_v2.reset()\n",
    "    model = DeepFM_PluginEmbedding(vocabulary_size=vocabulary_size, embedding_vec_size=embedding_vec_size, \n",
    "                                   embedding_type=embedding_type,\n",
    "                                   dropout_rate=[0.5] * 10, deep_layers=[1024] * 10,\n",
    "                                   initializer='uniform', gpus=gpus, batch_size=batch_size, \n",
    "                                   batch_size_eval=batch_size_eval,\n",
    "                                   slot_num=slot_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define training step\n",
    "@tf.function\n",
    "def _train_step(dense_batch, sparse_batch, y_batch, model, loss_fn, optimizer):\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_batch = tf.cast(y_batch, dtype=tf.float32)\n",
    "        logits = model(dense_batch, sparse_batch, training=True)\n",
    "        loss = loss_fn(y_batch, logits)\n",
    "        loss /= dense_batch.shape[0]\n",
    "    grads = tape.gradient(loss, model.trainable_weights)\n",
    "    optimizer.apply_gradients(zip(grads, model.trainable_weights))\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-01-30 07:52:25,346 begin to train\n",
      "2021-01-30 07:52:37,596 step: 100, loss: 0.0000278, elapsed time: 12.24864 seconds.\n",
      "2021-01-30 07:52:48,122 step: 200, loss: 0.0000301, elapsed time: 10.52632 seconds.\n",
      "2021-01-30 07:52:59,111 step: 300, loss: 0.0000292, elapsed time: 10.98891 seconds.\n",
      "2021-01-30 07:53:10,397 step: 400, loss: 0.0000298, elapsed time: 11.28664 seconds.\n",
      "2021-01-30 07:53:21,045 step: 500, loss: 0.0000308, elapsed time: 10.64784 seconds.\n",
      "2021-01-30 07:53:31,526 step: 600, loss: 0.0000298, elapsed time: 10.48030 seconds.\n",
      "2021-01-30 07:53:41,712 step: 700, loss: 0.0000298, elapsed time: 10.18635 seconds.\n",
      "2021-01-30 07:53:52,467 step: 800, loss: 0.0000304, elapsed time: 10.75503 seconds.\n",
      "2021-01-30 07:54:03,011 step: 900, loss: 0.0000299, elapsed time: 10.54400 seconds.\n",
      "2021-01-30 07:54:14,301 step: 1000, loss: 0.0000307, elapsed time: 11.28991 seconds.\n",
      "2021-01-30 07:54:25,194 step: 1100, loss: 0.0000286, elapsed time: 10.89364 seconds.\n",
      "2021-01-30 07:54:35,751 step: 1200, loss: 0.0000310, elapsed time: 10.55683 seconds.\n",
      "2021-01-30 07:54:46,374 step: 1300, loss: 0.0000318, elapsed time: 10.62237 seconds.\n",
      "2021-01-30 07:54:56,874 step: 1400, loss: 0.0000299, elapsed time: 10.50061 seconds.\n",
      "2021-01-30 07:55:07,540 step: 1500, loss: 0.0000308, elapsed time: 10.66558 seconds.\n",
      "2021-01-30 07:55:18,125 step: 1600, loss: 0.0000317, elapsed time: 10.58490 seconds.\n",
      "2021-01-30 07:55:28,575 step: 1700, loss: 0.0000298, elapsed time: 10.45029 seconds.\n",
      "2021-01-30 07:55:39,030 step: 1800, loss: 0.0000329, elapsed time: 10.45500 seconds.\n",
      "2021-01-30 07:55:49,386 step: 1900, loss: 0.0000326, elapsed time: 10.35561 seconds.\n",
      "2021-01-30 07:55:59,627 step: 2000, loss: 0.0000334, elapsed time: 10.24157 seconds.\n",
      "2021-01-30 07:56:09,869 step: 2100, loss: 0.0000335, elapsed time: 10.24189 seconds.\n",
      "2021-01-30 07:56:20,113 step: 2200, loss: 0.0000348, elapsed time: 10.24432 seconds.\n",
      "2021-01-30 07:56:24,112 Train End. Elapsed Time: 238.765 seconds.\n"
     ]
    }
   ],
   "source": [
    "# training loop\n",
    "logging.info(\"begin to train\")\n",
    "begin_time = time.time()\n",
    "display_begin = begin_time\n",
    "for step, datas in enumerate(dataset):\n",
    "    label, dense, others = datas[0], datas[1], datas[2:]\n",
    "    if \"Original\" == which_model:\n",
    "        sparse = others[-1]\n",
    "    elif \"Plugin\" == which_model:\n",
    "        sparse = others[0:2]\n",
    "        sparse = hugectr_tf_ops_v2.broadcast_then_convert_to_csr(model.get_embedding_name, \n",
    "                                                                row_indices=sparse[0], values=sparse[1], \n",
    "                                                                T=[tf.int32]*len(gpus))\n",
    "    train_loss = _train_step(dense, sparse, label, model, loss_fn, optimizer)\n",
    "    loss_value = train_loss.numpy()\n",
    "    \n",
    "    if (step % 100 == 0 and step != 0):\n",
    "        display_end = time.time()\n",
    "        logging.info(\"step: %d, loss: %.7f, elapsed time: %.5f seconds.\" %(step, loss_value, (display_end - display_begin)))\n",
    "        display_begin = display_end\n",
    "        \n",
    "end_time = time.time()\n",
    "logging.info(\"Train End. Elapsed Time: %.3f seconds.\" %(end_time - begin_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this configuration, `tf.data.Dataset` produces training data slowly, which makes the whole training process slow. Therefore, the training elapsed time for `Original` and `Plugin` are similar."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## API signature ##\n",
    "All embedding_plugin APIs are defined in [hugectr_tf_ops_v2.py](../tools/embedding_plugin/python/hugectr_tf_ops_v2.py).\n",
    "\n",
    "Embedding_plugin takes `COO (Coordinate)` format as input format when `fprop` is used. <br>\n",
    "In some cases, `fprop_experimental` can get better performance than `fprop`, but it is not stable. <br>\n",
    "If `fprop_experimental` is used, input data format should be `CSR (Compressed Sparse Row)`. \n",
    "\n",
    "For more detail about how to convert your input data to `CSR` or `COO` format, please refer to [samples/format_processing.py](../tools/embedding_plugin/samples/format_processing.py). <br>\n",
    "For more code samples, please refer to [samples/sample_with_fprop*.py](../tools/embedding_plugin/samples/sample_with_fprop.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>\n",
       "table {float:left}\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%html\n",
    "<style>\n",
    "table {float:left}\n",
    "</style>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  init(visible_gpus, seed=0, key_type='int64', value_type='float', batch_size=1, batch_size_eval=1)\n",
    "  ```\n",
    "  \n",
    "This function is used to create resource manager, which manages resources used by embedding_plugin.\n",
    "**IMPORTANT:** This function can only be called once. It must be called before any other embedding_plugin API is called. Currently, only key_type='int64', value_type='float' has been tested.\n",
    "\n",
    "\n",
    "| Args ||\n",
    "| :-----| :---- |\n",
    "| visible_gpus | list of integers, used to specify which gpus will be used by embedding_plugin. |\n",
    "| seed | integer, the initializer random seed for embedding_plugin. |\n",
    "| key_type| string, can be one of {'uint32', 'int64'}. Used to specify the input keys data type. |\n",
    "| value_type| string, can be one of {'float', 'half'}. Used to specify the data type of embedding_plugin forward result. |\n",
    "| batch_size| integer, batch_size used in training process. |\n",
    "| batch_size_eval| integer, batch_size used in evaluation process. |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  reset()\n",
    "  ```\n",
    "  \n",
    "This function is used to explicitly release the resources created by embedding plugin. You can call this function multiple times, and it will not throw any error."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  embedding_name = create_embedding(init_value, name_='hugectr_embedding', embedding_type='localized',\n",
    "                                     optimizer_type='Adam', max_vocabulary_size_per_gpu=1, slot_size_array=[],\n",
    "                                     opt_hparams=[0.001], update_type='Local', atomic_update=true, scaler=1.0,\n",
    "                                     slot_num=1, max_nnz=1, max_feature_num=1000, embedding_vec_size=1,\n",
    "                                     combiner='sum')\n",
    "  ```\n",
    "  \n",
    "| Args ||\n",
    "| :-----| :---- |\n",
    "|init_value| can be a `boolean` or a 2-D matrix with `dtype=tf.float32`. When it is `bool`, parameters will be randomly initialized. When it is a 2-D matrix with `dtype=tf.float32`, that matrix will be used to initialize embedding table parameters, and the matrix's row-index will be deemed to be key of the embedding table.|\n",
    "|name_|string, the name of this embedding layer. If `name_` is unique, then it will be used as the embedding layer name, otherwise, numerical suffix will be automatically added to `name_` to form an unique name for this embedding layer. |\n",
    "|embedding_type| string, can be one of {'localized', 'distributed'}. |\n",
    "| optimizer_type| string, can be one of {'Adam', 'MomentumSGD', 'Nesterov', 'SGD'}. | \n",
    "|max_vocabulary_size_per_gpu| integer, used to allocate GPU memory spaces for embedding layer.|\n",
    "|slot_size_array| list of integers, used to allocate GPU memory spaces precisely for embedding layer.|\n",
    "|opt_hparams| list of floats, used to specify hyper parameters for optimizer.<br>For `Adam`, `opt_hparams` must be a list of `[learning_rate, beta1, beta2, epsilon]`.<br>For `MomentumSGD`, `opt_hparams` must be a list of `[learning_rate, momentum_factor]`.<br>For `Nesterov`, `opt_hparams` must be a list of `[learning_rate, momentum_factor]`.<br>For `SGD`, `opt_hparams` must be a list of `[learning_rate]`.|\n",
    "|update_type| string, can be one of {'Local', 'Global', 'LazyGlobal'}. |\n",
    "|atomic_update| bool, only used in `SGD` optimizer. |\n",
    "|scaler| float, can be one of {1.0, 128.0, 256.0, 512.0, 1024.0}, used in `mixed_precission` training. |\n",
    "|slot_num| integer, how many slots (feature-fields) are unified in a single embedding layer. |\n",
    "|max_nnz| integer, the number of valid keys in a single slot.|\n",
    "|max_feature_num| integer, the number of valid keys in a single input sample.|\n",
    "|embedding_vec_size| integer, the embedding vector size of this embedding layer.|\n",
    "|combier|string, can be one of {'mean', 'sum'}. specify how to combine different embedding vector in the same slot.|\n",
    "\n",
    "|Returns||\n",
    "|:----| :---- |\n",
    "|embedding_name| tf.Tensor, dtype=tf.string. An unique name for this embedding layer.|\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  each_replica = broadcast_then_convert_to_csr(embedding_name, row_indices, values, T, is_training=True)\n",
    "  ```\n",
    "  \n",
    "This function take `COO` format as input on only one device. It will broadcast its input: `row_indices` and `values` to multiple GPUs, and then converting to `CSR` format will be conducted on each GPU independetely. This function **Must Not** be put inside of the scope of `MirroredStrategy` to avoid being called mutiple times in each iteration. After this function is called, `fprop` can be used to do the embedding forward propagation. \n",
    "\n",
    "For code sample, please refer to [samples/sample_with_fprop.py](../tools/embedding_plugin/samples/sample_with_fprop.py).\n",
    "\n",
    "|Args||\n",
    "|:----|:----|\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, use which embedding layer.|\n",
    "|row_indices| 1-D tf.Tensor with `dtype=tf.int64`. It denotes the `row indices` of the `COO` format input data for all GPUs.|\n",
    "|values| 1-D tf.Tensor with `dtype=tf.int64`. It denotes the `valid values` of the `COO` format input data for all GPUs.|\n",
    "| T | A list of `tf.int32`, whose length must be equal to the number of GPUs.|\n",
    "\n",
    "|Returns||\n",
    "|:----|:----|\n",
    "|each_replica| A list of tf.Tensor with `dtype=T`. Used as a placeholder to avoid this function being removed by `tf.Graph`.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  replica_forward = fprop(embedding_name, replica_id, to_each_replica, bp_trigger, is_training=True)\n",
    "  ```\n",
    "  \n",
    "This function can be used to do forward propagation for `Distributed` and `Localized` embedding layers. If multiple GPUs are used, this function must be put inside the scope of `MirroredStrategy` to do forward propagation on each GPU. If this function is used, `broadcast_then_convert_to_csr()` must be called before the calling of this function.\n",
    "\n",
    "For code sample, please refer to [samples/sample_with_fprop.py](../tools/embedding_plugin/samples/sample_with_fprop.py).\n",
    "\n",
    "|Args||\n",
    "|:----|:----|\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, use which embedding layer.|\n",
    "|replica_id|tf.Tensor with `dtype=tf.int32`, denotes the replica id in the group of GPUs. Can be obtained with `tf.distribute.get_replica_context().replica_id_in_sync_group`|\n",
    "|to_each_replica| Scalar tf.Tensor with `dtype=tf.int32`, the output from `broadcast_then_convert_to_csr()`|\n",
    "|bp_trigger| tf.Variable(dtype=tf.float32), used to automatically trigger back propagation of the embedding layer.|\n",
    "|is_training| `boolean`, specify whether use `training` resources or `evaluation` resources.|\n",
    "\n",
    "|Returns||\n",
    "|:----|:----|\n",
    "|replica_forward| tf.Tensor with `shape=[replica_batch_size, slot_num, embedding_vec_size]` for one GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  replica_forward = fprop_experimental(embedding_name, replica_id, row_offset, values, nnz, bp_trigger, is_training=True, input_buffer_reset=False)\n",
    "  ```\n",
    "  \n",
    "This function can be used to do forward propagation for `Distributed` and `Localized` embedding layers. If multiple GPUs are used, this function must be put inside the scope of `MirroredStrategy` to do forward propagation on each GPU. It only accept input data for one GPU, and produce the corresponding forward result for that GPU. And it takes `CSR` format as input.\n",
    "\n",
    "In some cases, this function can get better performance than `fprop`, but it is not stable. <br>\n",
    "For code sample, please refer to [samples/sample_with_fprop_experimental.py](../tools/embedding_plugin/samples/sample_with_fprop_experimental.py).\n",
    "\n",
    "|Args||\n",
    "|:----|:----|\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, use which embedding layer. |\n",
    "|replica_id| tf.Tensor with `dtype=tf.int32`, denotes the replica id in the group of GPUs. Can be obtained with `tf.distribute.get_replica_context().replica_id_in_sync_group`|\n",
    "|row_offset| 1-D tf.Tensor with `dtype=tf.int64`. It denotes the `row offset` of the `CSR` format input data for one GPU.|\n",
    "|values| 1-D tf.Tensor with `dtype=tf.int64`. It denotes the `valid value` of the `CSR` format input data for one GPU.|\n",
    "|nnz| Scalar tf.Tensor with `dtype=tf.int64`. It denotes the length of `values` of the `CSR` format input data for one GPU.|\n",
    "|bp_trigger|tf.Variable(dtype=tf.float32), used to automatically trigger back propagation of the embedding layer.|\n",
    "|is_training| `boolean`, specify whether use `training` resources or `evaluation` resources.|\n",
    "|input_buffer_reset| `boolean`, whether to reset the input buffer in each iteration.|\n",
    "\n",
    "|Returns||\n",
    "|:----|:----|\n",
    "|replica_forward| tf.Tensor with `shape=[replica_batch_size, slot_num, embedding_vec_size]` for one GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  save(embedding_name, save_name)\n",
    "  ```\n",
    "\n",
    "This function is used to save the `embedding_plugin` parameters in the file.\n",
    "\n",
    "|Args||\n",
    "|:----| :---- |\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, save which embedding layer's parameter to file.|\n",
    "|save_name| string, the name of saved parameters.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "---\n",
    "+ ```python\n",
    "  restore(embedding_name, file_name)\n",
    "  ```\n",
    "  \n",
    "This function is used to restore the `embedding_plugin` parameters from file.\n",
    "\n",
    "|Args||\n",
    "|:----| :---- |\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, restore parameters for which embedding layer.|\n",
    "|file_name| string, restore paramters from this file. |"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
